import sys
sys.path.extend(['.', '..'])
from c_ast_decl import *
from pycparser import c_generator
from c_ast_decl import ASTBuilder

builder = ASTBuilder()


def build_is_prime():
    """构建is_prime函数的AST"""
    # 函数参数

    n_param = create_variable_decl('n', 'int')

    
    # 函数声明
    func_decl = create_function_decl(
        'is_prime',
        'int',
        [n_param]
    )

    # 函数体内容
    body_stmts = []
    
    # if (n <= 1) return 0;
    body_stmts.append(
        c_ast.If(
            cond=c_ast.BinaryOp(
                '<=',
                c_ast.ID('n'),
                c_ast.Constant('int', '1')
            ),
            iftrue=c_ast.Return(c_ast.Constant('int', '0')),
            iffalse=None
        )
    )

    # if (n == 2) return 1;
    body_stmts.append(
        c_ast.If(
            cond=c_ast.BinaryOp(
                '==',
                c_ast.ID('n'),
                c_ast.Constant('int', '2')
            ),
            iftrue=c_ast.Return(c_ast.Constant('int', '1')),
            iffalse=None
        )
    )

    # if (n % 2 == 0) return 0;
    body_stmts.append(
        c_ast.If(
            cond=c_ast.BinaryOp(
                '==',
                c_ast.BinaryOp(
                    '%',
                    c_ast.ID('n'),
                    c_ast.Constant('int', '2')
                ),
                c_ast.Constant('int', '0')
            ),
            iftrue=c_ast.Return(c_ast.Constant('int', '0')),
            iffalse=None
        )
    )

    # for (int i = 3; i * i <= n; i += 2)
    for_loop = c_ast.For(
        init=create_variable_decl('i', 'int', init='3'),
        cond=c_ast.BinaryOp(
            '<=',
            c_ast.BinaryOp(
                '*',
                c_ast.ID('i'),
                c_ast.ID('i')
            ),
            c_ast.ID('n')
        ),
        next=c_ast.Assignment(
            '+=',
            c_ast.ID('i'),
            c_ast.Constant('int', '2')
        ),
        stmt=c_ast.Compound([
            c_ast.If(
                cond=c_ast.BinaryOp(
                    '==',
                    c_ast.BinaryOp(
                        '%',
                        c_ast.ID('n'),
                        c_ast.ID('i')
                    ),
                    c_ast.Constant('int', '0')
                ),
                iftrue=c_ast.Return(c_ast.Constant('int', '0')),
                iffalse=None
            )
        ])
    )
    body_stmts.append(for_loop)

    # return 1;
    body_stmts.append(c_ast.Return(c_ast.Constant('int', '1')))

    # 创建函数定义
    return create_function_def(
        func_decl = func_decl,
        body=c_ast.Compound(body_stmts)
    )

def build_main():
    """构建main函数的AST"""
    # 变量声明
    num_decl = create_variable_decl('num', 'int', init='17')

    # 函数调用
    func_call = c_ast.FuncCall(
        name=c_ast.ID('is_prime'),
        args=c_ast.ExprList([c_ast.ID('num')])
    )

    # if语句
    if_stmt = c_ast.If(
        cond=func_call,
        iftrue=c_ast.Compound([
            c_ast.FuncCall(
                name=c_ast.ID('printf'),
                args=c_ast.ExprList([
                    c_ast.Constant('string', r'"%d is prime\n"'),
                    c_ast.ID('num')
                ])
            )
        ]),
        iffalse=c_ast.Compound([
            c_ast.FuncCall(
                name=c_ast.ID('printf'),
                args=c_ast.ExprList([
                    c_ast.Constant('string', r'"%d is not prime\n"'),
                    c_ast.ID('num')
                ])
            )
        ])
    )

    # main函数定义
    return c_ast.FuncDef(
        decl=create_function_decl('main', 'int', []),
        param_decls=None,
        body=c_ast.Compound([
            num_decl,
            if_stmt,
            c_ast.Return(c_ast.Constant('int', '0'))
        ])
    )

# 构建完整AST
ast = c_ast.FileAST([
    build_is_prime(),
    build_main()
])

# 打印生成的C代码（需要自定义AST visitor）
# 这里只是示例结构，实际需要实现AST到代码的转换
# print("生成的AST结构：")
# print(ast)
generator = c_generator.CGenerator()
print(generator.visit(ast))
